
from search.graph import *
from search.astar import *
import osm_utils2
import wayzUtils
import problem_agent
import sys
import gc
from itertools import izip
class NodeExtend(Node):
    def __init__(self, state, parent=None, action=None, path_cost=0):
        self.ParentActionDic = {} # add dic parent junction key-> cost 
        Node.__init__(self, state, parent, action, path_cost)
    def expand(self):
        '''
        Return a list of nodes reachable from this node.
        '''
        def path_cost(action):
            return self.path_cost + action.cost

        successors = self.state.getSuccessors()
        return [NodeExtend(next, self, act, path_cost(act))
                        for (act, next) in successors.items()]
class SpeedUpdatesSearch(GraphSearch):
    
    def find(self, problem_state, heuristic=None):
        def updateParentandCost(NodeExtend,closed_states):
            tempCost = -1
            tempParent = -1
            for par in NodeExtend.ParentActionDic:
                if NodeExtend.ParentActionDic[par] < tempCost or tempCost < 0:
                    tempCost = NodeExtend.ParentActionDic[par]
                    tempParent = par
            if tempParent != -1:
                closed_states[NodeExtend.state.currentJunctionKey].path_cost = tempCost
                closed_states[NodeExtend.state.currentJunctionKey].parent = closed_states[tempParent]
        needToUpdated = False 
        dicOfOpenState = {} 
        open_states = self.container_generator()
        reopen_states = self.container_generator()
        reopened = False
        closed_states = {}
        #
        list_of_links = []
        #
        SourceNode = NodeExtend(problem_state)

        open_states.append(SourceNode)
        
        #
        problem_state.roadMap.ZeroChangeCounter()
        #  
        num =0
        while open_states and len(open_states) > 0:
            
            
            #
            LinksToUpdate = SourceNode.state.roadMap.GetSpeedUpdates(( list_of_links))
#            print "size of LinksToUpdate",len(LinksToUpdate)
            for link in LinksToUpdate:#links to update
                needToUpdated = True
                tempStartJunctionKey = link[0][0]#start of link
                tempEndJunctionKey = link[0][1]#end of link
                tempJunction = SourceNode.state.roadMap.GetJunction(tempStartJunctionKey)#start junction data 
                for i in tempJunction.links:#for all links of start junction
                    if i.target == tempEndJunctionKey:#search for correct link start->end
                        i.speed=int(i.speed*(1 + link[1])) #update link
                        if i.speed > 120:
                            i.speed = 120
                        if i.speed < 0:
                            i.speed = 0
                        temp_node = closed_states[tempStartJunctionKey].expand()#recreate node with right cost(brothers are made as well with no use)
                        for newNode in temp_node:#go on new nodes
                            if newNode.state.currentJunctionKey == i.target:#search for right new node
                                originalCost = closed_states[tempEndJunctionKey].ParentActionDic[tempStartJunctionKey]
                                if newNode.path_cost < originalCost:
                                    closed_states[tempEndJunctionKey].ParentActionDic[tempStartJunctionKey] = newNode.path_cost
                                    updateParentandCost(closed_states[tempEndJunctionKey],closed_states)
                                    open_states.append(newNode)
                                else:
                                    if newNode.path_cost != originalCost:# change in cost
                                            ####update father and cost#####
                                        closed_states[tempEndJunctionKey].ParentActionDic[tempStartJunctionKey] = newNode.path_cost
                                        updateParentandCost(closed_states[tempEndJunctionKey],closed_states)
                                        reopen_states.append(newNode)
                                        reopened = True
            if not needToUpdated:
                allOpened = (open_states.q)
                i = iter(allOpened)
                dicOfOpenState = dict(izip(i, i)) 
                needToUpdated = False
                                   
            num +=1
            SourceNode.parent = None
            if reopened and len(reopen_states) > 0:
                node = reopen_states.pop()
                if node.state.currentJunctionKey not in closed_states:
                    continue
            else:
                reopened = False
                node = open_states.pop()
            if node.depth > self.max_depth:
                continue
            
            if num % 100000 == 0: #once in a while clean memory
                gc.collect()
                
            if num > 1868558:
                return None
            if num > 1 and not reopened:#add new roads to list
                list_of_links.append([node.parent.state.currentJunctionKey, node.state.currentJunctionKey])
            SourceNode.parent = None#make sure Source Node father is none
            if node.state.isGoal(): 
                del closed_states
                gc.collect()
                return node.getPathActions()
            if (node.state.currentJunctionKey not in closed_states):
                closed_states[node.state.currentJunctionKey] = node
                closed_states[node.state.currentJunctionKey].path_cost = node.path_cost
                if num > 1:
                    closed_states[node.state.currentJunctionKey].ParentActionDic[node.parent.state.currentJunctionKey] = node.path_cost
                if reopened:
                    #if reopen add to reopen new nodes
                    succ = node.expand()
                    for new_node in succ:
                        if new_node not in dicOfOpenState and new_node in closed_states:
                            reopen_states.append(closed_states[new_node.state.currentJunctionKey])
                else:
                    SourceNode.parent = None
                    open_states.extend(node.expand())
                
            else:
                #update node dic
                closed_states[node.state.currentJunctionKey].ParentActionDic[node.parent.state.currentJunctionKey] = node.path_cost
                if closed_states[node.state.currentJunctionKey].path_cost > node.path_cost:
                    closed_states[node.state.currentJunctionKey].path_cost = node.path_cost
                    closed_states[node.state.currentJunctionKey].parent = node.parent 
                    if reopened:
                        #if reopen add to reopen new nodes
                        succ = node.expand()
                        for new_node in succ:
                            if new_node not in dicOfOpenState:
                                reopen_states.append(closed_states[new_node.state.currentJunctionKey])
                    else:
                        SourceNode.parent = None
                        open_states.extend(node.expand())
            SourceNode.parent = None               
        return None
    
class SpeedUpdatesSearchNaive(GraphSearch):
    
    def find(self, problem_state, heuristic=None):
        open_states = self.container_generator()
        closed_states = {}
        #
        list_of_links = []
        #
        open_states.append(Node(problem_state))
        
        #
        problem_state.roadMap.ZeroChangeCounter()
        #
        num = 0
        while open_states and len(open_states) > 0:
            num += 1
            node = open_states.pop()
                       
            if node.depth > self.max_depth:
                continue
            #update roads
            LinksToUpdate = node.state.roadMap.GetSpeedUpdates(( list_of_links))
            for link in LinksToUpdate:
                tempJunctionKey = link[0][0]
                tempLinkKey = link[0][1]
                tempJunction = node.state.roadMap.GetJunction(tempJunctionKey)
                for i in tempJunction.links:
                    if i == tempLinkKey:
                        i.speed=int(i.speed*(1 + link[1])) 
                if len(tempJunction.links) > 0:
                    open_states = self.container_generator()
                    closed_states = {}
                    num = 0
                    open_states.append(Node(problem_state))
                    gc.collect()
                    
            #
            if node.state.isGoal(): 
                return node.getPathActions()
            if num > 1 and len(node.getPathActions())>0:
                list_of_links.append([node.parent.state.currentJunctionKey, node.state.currentJunctionKey])
            if (node.state not in closed_states) or (node.path_cost < closed_states[node.state].path_cost):
                closed_states[node.state] = node
                #
                open_states.extend(node.expand())
                
        return None
    


class SpeedUpdatesAStar(AStar):
    def find(self, problem_state, heuristic):
        '''
        A* search is best-first graph search with f(n) = g(n)+h(n).
        You need to specify the h function when you call astar_search.
        Uses the pathmax trick: f(n) = max(f(n), g(n)+h(n)).

        @param problem_state: The initial state to start the search from.
        @param heuristic: A heuristic function that scores the Problem states.
        '''
        # This is the node evaluation function for the given heuristic.
        def evaluator(node):
            return node.path_cost + heuristic.evaluate(node.state)
    
        # This is a generator for the PriorityQueue we need.
        def queue_generator():
            return PriorityQueue(evaluator)
        # Use a graph search with a minimum priority queue to conduct the search.
        search = SpeedUpdatesSearch(queue_generator, self.max_depth)
        return search.find(problem_state)
    
class WayzProblemAgentUpdatesAStar(problem_agent.ProblemAgent):
    def solve(self, problem_state, time_limit = infinity):
        aStarSearch = SpeedUpdatesAStar(time_limit)
        myHeuristic = wayzUtils.CompositeHeuristic(problem_state.carProfile)
        return aStarSearch.find(problem_state,myHeuristic)
    
    
class SpeedUpdatesAStarNaive(AStar):
    def find(self, problem_state, heuristic):
        '''
        A* search is best-first graph search with f(n) = g(n)+h(n).
        You need to specify the h function when you call astar_search.
        Uses the pathmax trick: f(n) = max(f(n), g(n)+h(n)).

        @param problem_state: The initial state to start the search from.
        @param heuristic: A heuristic function that scores the Problem states.
        '''
        # This is the node evaluation function for the given heuristic.
        def evaluator(node):
            return node.path_cost + heuristic.evaluate(node.state)
    
        # This is a generator for the PriorityQueue we need.
        def queue_generator():
            return PriorityQueue(evaluator)
        # Use a graph search with a minimum priority queue to conduct the search.
        search = SpeedUpdatesSearchNaive(queue_generator, self.max_depth)
        return search.find(problem_state)
    
class WayzProblemAgentUpdatesAStarNaive(problem_agent.ProblemAgent):
    def solve(self, problem_state, time_limit = infinity):
        aStarSearch = SpeedUpdatesAStarNaive(time_limit)
        myHeuristic = wayzUtils.CompositeHeuristic(problem_state.carProfile)
        return aStarSearch.find(problem_state,myHeuristic)